# coding=utf-8
from expr_optimal import schedule_optimal
from expr_greedy import schedule_greedy
from expr_swap import schedule_swap
from chiplet import QPU
from dag import Circuit_DAG

# 算法类型 bv / qaoa / qft / vqa
# 量子比特数 20, 40, 60, 80, 100, 120, 140, 160, 180, 200

circuit_type = ["bv", "qaoa", "qft", "vqa"]
qubit_num = [20, 60, 100, 140, 180]
chiplet_type = ["square"]

chiplet_size = [
    {"square": 8, "hexagon": 5, "heavy_square": 5, "heavy_hexagon": 3},
    {"square": 8, "hexagon": 5, "heavy_square": 5, "heavy_hexagon": 3},
    {"square": 8, "hexagon": 5, "heavy_square": 5, "heavy_hexagon": 3},
    {"square": 8, "hexagon": 5, "heavy_square": 5, "heavy_hexagon": 3},
    {"square": 8, "hexagon": 5, "heavy_square": 5, "heavy_hexagon": 3},
]

chiplet_array = [
    {"square": [3, 3], "hexagon": [3, 3],
        "heavy_square": [3, 3], "heavy_hexagon": [3, 3]},
    {"square": [3, 3], "hexagon": [3, 3],
        "heavy_square": [3, 3], "heavy_hexagon": [3, 3]},
    {"square": [3, 3], "hexagon": [3, 3],
        "heavy_square": [3, 3], "heavy_hexagon": [3, 3]},
    {"square": [3, 3], "hexagon": [3, 3],
        "heavy_square": [3, 3], "heavy_hexagon": [3, 3]},
    {"square": [3, 3], "hexagon": [3, 3],
        "heavy_square": [3, 3], "heavy_hexagon": [3, 3]},
]


def task(num_idx, num_item, circuit_type_item, chiplet_type_item):
    # 这里是原来嵌套循环中的处理逻辑
    # QASM文件
    qasm_file = f"circuits/{circuit_type_item}_{num_item}.qasm"
    # 电路的DAG图
    circuit_dag = Circuit_DAG(qasm_file)
    # Chip内大小
    inner_chip_size = chiplet_size[num_idx][chiplet_type_item]
    # 芯片矩阵大小
    inter_chip_array = chiplet_array[num_idx][chiplet_type_item]
    total_count = {
        "swap": [],
        "greedy": [],
        "optimal": [],
    }
    qpu = QPU(chip_type=chiplet_type_item, inner_chip_size=inner_chip_size, inter_chip_array=inter_chip_array,
              circuit_dag=circuit_dag, just_one_map=(circuit_type == "vqa" or circuit_type == "qft"))
    # for func_idx, func in enumerate([schedule_swap, schedule_greedy, schedule_optimal]):
    for func_idx, func in enumerate([schedule_swap, schedule_greedy]):
        max_depth, ancilla_num = func(qpu)
        arr = ["swap", "greedy", "optimal"]
        total_count[arr[func_idx]].append(max_depth)
        total_count[arr[func_idx]].append(ancilla_num)

    output_file_path = "output_qubit_number_simple.txt"
    with open(output_file_path, 'a') as file:
        file.write(
            f"{num_item} {circuit_type_item} {chiplet_type_item} {inner_chip_size} {inter_chip_array} {total_count}\n")
    print(num_item, circuit_type_item, chiplet_type_item,
          inner_chip_size, inter_chip_array, total_count)


for num_idx, num_item in enumerate(qubit_num):  # 20, 60, 100, 140, 180
    for circuit_type_item in circuit_type:  # "bv", "qaoa", "qft", "vqa"
        for chiplet_type_item in chiplet_type:  # "square"
            task(num_idx, num_item, circuit_type_item, chiplet_type_item)
